#!/bin/bash
# This script may be used to run fio uNVMe benchmarks.
# It determines which fio engine to use based on the specified device name.
#
# Usage examples:
#   % unvme-benchmark /dev/nvme0n1  # using kernel space NVMe driver
#   % unvme-benchmark 05:00.0       # using user space uNVMe driver

PROG=$(basename $0)
USAGE="Usage: ${PROG} DEVICE_NAME"

cd $(dirname $0)
: ${OUTDIR=${PWD}/out}

FIOTEXT="[global]
norandommap=1
thread=1
group_reporting=1
direct=1
verify=0
ioengine=IOENGINE
time_based=1
ramp_time=RAMPTIME
runtime=RUNTIME
bs=BLOCKSIZE
filename=FILENAME

[test]
rw=RW
numjobs=NUMJOBS
iodepth=IODEPTH"

: ${FIODIR=$(grep '^FIODIR' ../Makefile.def | sed 's/.*=\s*//')}
: ${RW="randwrite randread"}
: ${NUMJOBS="01 08 16"}
: ${IODEPTH="01 08 16 32"}
: ${RAMPTIME=60}
: ${RUNTIME=120}
: ${BLOCKSIZE=4096}
: ${IOENGINE=""}

FIO=${FIODIR}/fio

[ $# -lt 1 ] && echo ${USAGE} && exit 1
[ ${EUID} -ne 0 ] && echo "ERROR: ${PROG} must be run as root" && exit 1
[ ! -x "${FIO}" ] && echo "ERROR: Missing ${FIO}" && exit 1

for i in $@; do
    case $i in
    /dev/nvme*)
        if [ ! -b $i ]; then
            echo "Unknown NVMe device: $i"
            exit 1
        fi
        FILENAME=$i
        [ -z "${IOENGINE}" ] && IOENGINE=libaio
        DRIVER="nvme"
        ;;

    [0-9A-Fa-f][0-9A-Fa-f]:[0-9A-Fa-f][0-9A-Fa-f].[0-9A-Fa-f]*)
        pcidev=$(echo $i | cut -d/ -f1)
        if [ -z "$(lspci -n | grep ${pcidev}\ 0108:)" ]; then
            echo "Device ${pcidev} is not SSD"
            exit 1
        fi
        if [ ! -e /sys/bus/pci/drivers/vfio-pci/0000:${pcidev} ]; then
            echo "Device ${pcidev} is not uNVMe enabled"
            exit 1
        fi
        FILENAME=$(echo $i | tr :/ .)
        [ -z "${IOENGINE}" ] && IOENGINE=$(readlink -f ../ioengine/unvme_fio)
        DRIVER="unvme"
        [ -n "$(pgrep unvmed)" ] && DRIVER="unvmed"
        ;;

    *)
        echo ${USAGE}
        exit 1
        ;;
    esac
done

mkdir -p ${OUTDIR}
[ $? -ne 0 ] && exit 1


#
# Echo and execute a command.
#
excmd() {
    echo "# $* ($(date '+%a %D %r'))"
    eval $*
    [ $? -ne 0 ] && exit 1
    echo
    #read -p "Press <Enter> to continue..."
}


#
# Start test.
#
for rw in ${RW}; do
    for qc in ${NUMJOBS}; do
        for qd in ${IODEPTH}; do
            OUTNAME=${OUTDIR}/${DRIVER}-${rw}-${qc}-${qd}
            OUTFILE=${OUTNAME}.out
            FIOFILE=${OUTNAME}.fio

            if [ -e ${OUTFILE} ]; then
                if [ -n "$(grep 'Run status' ${OUTFILE})" ]; then
                    echo -e "Skip existing ${OUTFILE}\n"
                    continue
                else
                    excmd rm -f ${OUTFILE}
                fi
            fi

            echo "${FIOTEXT}" | sed -e "s?IOENGINE?${IOENGINE}?g;s?RAMPTIME?${RAMPTIME}?g;s?RUNTIME?${RUNTIME}?g;s?BLOCKSIZE?${BLOCKSIZE}?g;s?FILENAME?${FILENAME}?g;s?RW?${rw}?g;s?NUMJOBS?${qc}?g;s?IODEPTH?${qd}?g" > ${FIOFILE}

            echo "========"
            excmd /bin/uname -a | tee -a ${OUTFILE}
            excmd cat ${FIOFILE} | tee -a ${OUTFILE}
            echo -e "\n========\n" >> ${OUTFILE}
            excmd /usr/bin/free -h | tee -a ${OUTFILE}
            echo -e "\n========\n" >> ${OUTFILE}

            excmd ${FIO} ${FIOFILE} 2>&1 | tee -a ${OUTFILE}

            if [ -n "$(grep 'Run status' ${OUTFILE})" ]; then
                echo -e "\n========\n" >> ${OUTFILE}
                excmd /usr/bin/free -h | tee -a ${OUTFILE}
                rm ${FIOFILE}
            else
                echo "### ${PROG} terminated on error ###"
                exit 1
            fi
        done
    done
done

